import matplotlib.pyplot as plt
import numpy as np

# 图2.1 强化学习智能体与环境交互示意图
fig, ax = plt.subplots(1, 1, figsize=(8, 5))
fig.suptitle('图2.1 强化学习智能体-环境交互回路', fontproperties='SimHei', fontsize=14)

# 绘制智能体和环境框图
agent = plt.Rectangle((0.2, 0.3), 0.2, 0.4, fc='lightblue', ec='black')
env = plt.Rectangle((0.6, 0.3), 0.2, 0.4, fc='lightgreen', ec='black')
ax.add_patch(agent)
ax.add_patch(env)

ax.text(0.3, 0.7, '智能体\n(Agent)', ha='center', va='center', fontproperties='SimHei')
ax.text(0.7, 0.7, '环境\n(Environment)', ha='center', va='center', fontproperties='SimHei')

# 绘制箭头表示交互
ax.arrow(0.4, 0.5, 0.18, 0, head_width=0.02, head_length=0.02, fc='r', ec='r', label='动作 $a_t$')
ax.arrow(0.6, 0.5, -0.18, 0, head_width=0.02, head_length=0.02, fc='g', ec='g', label='状态 $s_t$, 奖励 $r_t$')

ax.text(0.5, 0.53, '动作 $a_t$', ha='center', va='bottom', color='red', fontsize=12)
ax.text(0.5, 0.47, '状态 $s_{t+1}$\n奖励 $r_{t+1}$', ha='center', va='top', color='green', fontsize=10)

ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.axis('off')
ax.legend(loc='upper center', bbox_to_anchor=(0.5, 0.15))
plt.tight_layout()
plt.show()